/*
* Copyright (C) Microsoft. All rights reserved.  
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not 
* use this file except in compliance with the License.  You may obtain a copy 
* of the License at http://www.apache.org/licenses/LICENSE-2.0  
* 
* THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED 
* WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE, 
* MERCHANTABLITY OR NON-INFRINGEMENT. 
*
* See the Apache License, Version 2.0 for specific language governing 
* permissions and limitations under the License.
*
*/
//
// Copyright (c) Microsoft Corporation.  All rights reserved.
//
//
// Use of this source code is subject to the terms of the Microsoft shared
// source or premium shared source license agreement under which you licensed
// this source code. If you did not accept the terms of the license agreement,
// you are not authorized to use this source code. For the terms of the license,
// please see the license agreement between you and Microsoft or, if applicable,
// see the SOURCE.RTF on your install media or the root of your tools installation.
// THE SOURCE CODE IS PROVIDED "AS IS", WITH NO WARRANTIES OR INDEMNITIES.
//

#include "pch.h"
#include "dxtexhelpers.h"

using namespace DirectX;

namespace VisualStudioGraphics
{
    
    HRESULT
    CreateTextureFromSystemMemory(
        __in ID3D11Device* d3d11Device,
        __in DXGI_FORMAT format,
        __in GxMipmapLoadEnum mipmapLoading,
        __in D3D11_USAGE usage,
        __in uint32 bindFlags,
        __in uint32 cpuAccessFlags,
        __in const ScratchImage& image,
        __deref_out ID3D11Resource** d3d11Resource
        )
    {
        E_INVALIDARG_Return(format != DXGI_FORMAT_UNKNOWN);
        E_POINTER_Return(d3d11Device);
        E_POINTER_Return(d3d11Resource);

        TexMetadata metadata = image.GetMetadata();
        const Image* images = image.GetImages();
        uint32 numImages = image.GetImageCount();
        E_INVALIDARG_Return(metadata.mipLevels > 0 && metadata.arraySize > 0);
        E_POINTER_Return(images);
        E_INVALIDARG_Return(numImages > 0);
        
        TRACEHR hr = S_OK;
        
        ScratchImage convertedImage;
        if (format != DXGI_FORMAT_FROM_FILE && format != metadata.format)
        {
            hr = ConvertPixelFormat(format, image, &convertedImage);
            if (SUCCEEDED(hr))
            {
                metadata = convertedImage.GetMetadata();
                images = convertedImage.GetImages();
                numImages = convertedImage.GetImageCount();
            }
        }

        ScratchImage autoGeneratedMips;
        if (SUCCEEDED(hr))
        {
            if (mipmapLoading == GxMipmapLoadNoMipmaps)
            {
                // only load the top level image
                metadata.mipLevels = 1;
            }
            else if (mipmapLoading == GxMipmapLoadForceMipmaps)
            {
                hr = GenerateFullMipChain(images[0], &autoGeneratedMips);
                if (SUCCEEDED(hr))
                {
                    metadata = autoGeneratedMips.GetMetadata();
                    images = autoGeneratedMips.GetImages();
                    numImages = autoGeneratedMips.GetImageCount();
                }
            }
        }
        
        // fill out subresource array
        vector<D3D11_SUBRESOURCE_DATA> initData(metadata.mipLevels * metadata.arraySize);
        if (SUCCEEDED(hr))
        {
            if (metadata.dimension == TEX_DIMENSION_TEXTURE3D)
            {
                hr = HrFromBool(metadata.depth > 0, E_INVALIDARG);
                if (SUCCEEDED(hr))
                {
                    // Direct3D 11 doesn't support arrays of 3D textures
                    if (metadata.arraySize > 1)
                    {
                        hr = HRESULT_FROM_WIN32(ERROR_NOT_SUPPORTED);
                    }
                }
                
                uint32 depth = metadata.depth;
                uint32 idx = 0;
                for (uint32 level = 0; SUCCEEDED(hr) && level < metadata.mipLevels; level++)
                {
                    uint32 index = metadata.ComputeIndex(level, 0, 0);
                    hr = HrFromBool(index < numImages);
                    if (SUCCEEDED(hr))
                    {
                        const Image& img = images[index];
                        if (img.format != metadata.format)
                        {
                            hr = E_FAIL;
                        }
                        else if (img.pixels == NULL)
                        {
                            hr = E_POINTER;
                        }
                        else
                        {
                            // Verify pixels in image 1 .. (depth-1) are exactly image->slicePitch apart
                            // For 3D textures, this relies on all slices of the same miplevel being continous in memory
                            // (this is how ScratchImage lays them out), which is why we just give the 0th slice to Direct3D 11
                            const uint8_t* pSlice = img.pixels + img.slicePitch;
                            for (uint32 slice = 1; SUCCEEDED(hr) && slice < depth; slice++)
                            {
                                uint32 tindex = metadata.ComputeIndex(level, 0, slice);
                                hr = HrFromBool(tindex < numImages);
                                if (SUCCEEDED(hr))
                                {
                                    const Image& timg = images[tindex];
                                    if (timg.pixels == NULL)
                                    {
                                        hr = E_POINTER;
                                    }
                                    else if (timg.pixels != pSlice || timg.format != metadata.format || timg.rowPitch != img.rowPitch || timg.slicePitch != img.slicePitch)
                                    {
                                        hr = E_FAIL;
                                    }
                                    else 
                                    {
                                        pSlice = timg.pixels + img.slicePitch;
                                    }
                                }
                            }

                            if (SUCCEEDED(hr))
                            {
                                VSGRAPHICSASSERT(idx < (metadata.mipLevels * metadata.arraySize));
                                initData[idx].pSysMem = img.pixels;
                                initData[idx].SysMemPitch = static_cast<DWORD>(img.rowPitch);
                                initData[idx].SysMemSlicePitch = static_cast<DWORD>(img.slicePitch);
                                ++idx;
                                
                                if (depth > 1)
                                {
                                    depth >>= 1;
                                }
                            }
                        }
                    }
                }
            }
            else
            {
                //--- 1D or 2D texture case ---------------------------------------------------
                
                uint32 idx = 0;
                for (uint32 item = 0; SUCCEEDED(hr) && item < metadata.arraySize; item++)
                {
                    for (uint32 level = 0; SUCCEEDED(hr) && level < metadata.mipLevels; level++)
                    {
                        uint32 index = metadata.ComputeIndex(level, item, 0);
                        hr = HrFromBool(index < numImages);
                        if (SUCCEEDED(hr))
                        {
                            const Image& img = images[index];
                            if (img.format != metadata.format)
                            {
                                hr = E_FAIL;
                            }
                            else if (img.pixels == NULL)
                            {
                                hr = E_POINTER;
                            }
                            else
                            {
                                VSGRAPHICSASSERT(idx < (metadata.mipLevels * metadata.arraySize));
                                initData[idx].pSysMem = img.pixels;
                                initData[idx].SysMemPitch = static_cast<DWORD>( img.rowPitch );
                                initData[idx].SysMemSlicePitch = static_cast<DWORD>( img.slicePitch );
                                ++idx;
                            }
                        }
                    }
                }
            }
        }
        
        if (SUCCEEDED(hr))
        {
            switch (metadata.dimension)
            {
            case TEX_DIMENSION_TEXTURE1D:
                {
                    D3D11_TEXTURE1D_DESC desc;
                    desc.Width = metadata.width;
                    desc.MipLevels = metadata.mipLevels;
                    desc.ArraySize = metadata.arraySize;
                    desc.Format = metadata.format;
                    desc.Usage = usage;
                    desc.BindFlags = bindFlags;
                    desc.CPUAccessFlags = cpuAccessFlags;
                    desc.MiscFlags = 0;
                    
                    hr = d3d11Device->CreateTexture1D(&desc, initData.data(), reinterpret_cast<ID3D11Texture1D**>(d3d11Resource));
                }
                break;
            
            case TEX_DIMENSION_TEXTURE2D:
                {
                    D3D11_TEXTURE2D_DESC desc;
                    desc.Width = metadata.width;
                    desc.Height = metadata.height;
                    desc.MipLevels = metadata.mipLevels;
                    desc.ArraySize = metadata.arraySize;
                    desc.Format = metadata.format;
                    desc.SampleDesc.Count = 1;
                    desc.SampleDesc.Quality = 0;
                    desc.Usage = usage;
                    desc.BindFlags = bindFlags;
                    desc.CPUAccessFlags = cpuAccessFlags;
                    desc.MiscFlags = (metadata.miscFlags & TEX_MISC_TEXTURECUBE) ? D3D11_RESOURCE_MISC_TEXTURECUBE : 0;

                    hr = d3d11Device->CreateTexture2D(&desc, initData.data(), reinterpret_cast<ID3D11Texture2D**>(d3d11Resource));
                }
                break;
            
            case TEX_DIMENSION_TEXTURE3D:
                {
                    D3D11_TEXTURE3D_DESC desc;
                    desc.Width = metadata.width;
                    desc.Height = metadata.height;
                    desc.Depth = metadata.depth;
                    desc.MipLevels = metadata.mipLevels;
                    desc.Format = metadata.format;
                    desc.Usage = usage;
                    desc.BindFlags = bindFlags;
                    desc.CPUAccessFlags = cpuAccessFlags;
                    desc.MiscFlags = 0;

                    hr = d3d11Device->CreateTexture3D(&desc, initData.data(), reinterpret_cast<ID3D11Texture3D**>(d3d11Resource));
                }
                break;

            default:
                hr = E_FAIL;
                break;
            }
        }
        
        return hr;
    }
    
    inline uint32
    ComputeScanlines(
        __in DXGI_FORMAT format, 
        __in uint32 height
        )
    {
        switch (format)
        {
        case DXGI_FORMAT_BC1_TYPELESS:
        case DXGI_FORMAT_BC1_UNORM:
        case DXGI_FORMAT_BC1_UNORM_SRGB:
        case DXGI_FORMAT_BC2_TYPELESS:
        case DXGI_FORMAT_BC2_UNORM:
        case DXGI_FORMAT_BC2_UNORM_SRGB:
        case DXGI_FORMAT_BC3_TYPELESS:
        case DXGI_FORMAT_BC3_UNORM:
        case DXGI_FORMAT_BC3_UNORM_SRGB:
        case DXGI_FORMAT_BC4_TYPELESS:
        case DXGI_FORMAT_BC4_UNORM:
        case DXGI_FORMAT_BC4_SNORM:
        case DXGI_FORMAT_BC5_TYPELESS:
        case DXGI_FORMAT_BC5_UNORM:
        case DXGI_FORMAT_BC5_SNORM:
        case DXGI_FORMAT_BC6H_TYPELESS:
        case DXGI_FORMAT_BC6H_UF16:
        case DXGI_FORMAT_BC6H_SF16:
        case DXGI_FORMAT_BC7_TYPELESS:
        case DXGI_FORMAT_BC7_UNORM:
        case DXGI_FORMAT_BC7_UNORM_SRGB:
            return max(1, (height + 3) / 4);

        default:
            return height;
        }
    }
    
    static HRESULT
    _Capture(
        __in ID3D11DeviceContext* d3d11Context,
        __in ID3D11Resource* d3d11Resource,
        __in const TexMetadata& metadata,
        __deref_out ScratchImage* image
        )
    {
        E_POINTER_Return(d3d11Context);
        E_POINTER_Return(d3d11Resource);
        E_POINTER_Return(image);
        
        TRACEHR hr = S_OK;

        if (metadata.dimension == TEX_DIMENSION_TEXTURE3D)
        {
            VSGRAPHICSASSERT(metadata.arraySize == 1);
            
            uint32 height = metadata.height;
            uint32 depth = metadata.depth;
            for (uint32 level = 0; SUCCEEDED(hr) && level < metadata.mipLevels; level++)
            {
                D3D11_MAPPED_SUBRESOURCE mapped;
                uint32 dindex = D3D11CalcSubresource(level, 0, metadata.mipLevels);
                const uint8_t* pslice = NULL;
                uint32 lines = 0;
                bool isMapped = false;

                hr = d3d11Context->Map(d3d11Resource, dindex, D3D11_MAP_READ, 0, &mapped);
                if (SUCCEEDED(hr))
                {
                    isMapped = true;

                    pslice = reinterpret_cast<const uint8_t*>(mapped.pData);
                    hr = HrFromPointer(pslice);
                    if (SUCCEEDED(hr))
                    {
                        lines = ComputeScanlines(metadata.format, height);
                    }
                }

                for (uint32 slice = 0; slice < depth; slice++)
                {
                    const Image* img = image->GetImage(level, 0, slice);
                    hr = HrFromPointer(img);
                    if (SUCCEEDED(hr))
                    {
                        hr = HrFromPointer(img->pixels);
                    }

                    if (SUCCEEDED(hr))
                    {
                        const uint8_t* sptr = pslice;
                        uint8_t* dptr = img->pixels;
                        for (uint32 h = 0; h < lines; h++)
                        {
                            uint32 msize = min(img->rowPitch, mapped.RowPitch);
                            memcpy_s(dptr, img->rowPitch, sptr, msize);
                            sptr += mapped.RowPitch;
                            dptr += img->rowPitch;
                        }
                        
                        pslice += mapped.DepthPitch;
                    }
                }

                if (isMapped)
                {
                    d3d11Context->Unmap(d3d11Resource, dindex);
                }
                
                if (height > 1)
                {
                    height >>= 1;
                }
                if (depth > 1)
                {
                    depth >>= 1;
                }
            }
        }
        else
        {
            // 1D or 2D texture
            VSGRAPHICSASSERT(metadata.depth == 1);
            for (uint32 item = 0; SUCCEEDED(hr) && item < metadata.arraySize; item++)
            {
                uint32 height = metadata.height;
                for (uint32 level = 0; SUCCEEDED(hr) && level < metadata.mipLevels; level++)
                {
                    uint32 lines = 0;
                    uint32 dindex = D3D11CalcSubresource(level, item, metadata.mipLevels);
                    D3D11_MAPPED_SUBRESOURCE mapped;
                    bool isMapped = false;

                    hr = d3d11Context->Map(d3d11Resource, dindex, D3D11_MAP_READ, 0, &mapped);
                    if (SUCCEEDED(hr))
                    {
                        isMapped = true;

                        const Image* img = image->GetImage(level, item, 0);
                        hr = HrFromPointer(img);
                        if (SUCCEEDED(hr))
                        {
                            hr = HrFromPointer(img->pixels);
                            if (SUCCEEDED(hr))
                            {
                                lines = ComputeScanlines(metadata.format, height);
                            }
                        }
                        
                        if (SUCCEEDED(hr))
                        {
                            const uint8_t* sptr = reinterpret_cast<const uint8_t*>(mapped.pData);
                            uint8_t* dptr = img->pixels;
                            for (uint32 h = 0; h < lines; h++)
                            {
                                uint32 msize = min(img->rowPitch, mapped.RowPitch);
                                memcpy_s(dptr, img->rowPitch, sptr, msize);
                                sptr += mapped.RowPitch;
                                dptr += img->rowPitch;
                            }
                        }
                    }

                    if (isMapped)
                    {
                        d3d11Context->Unmap(d3d11Resource, dindex);
                    }
                    
                    if (height > 1)
                    {
                        height >>= 1;
                    }
                }
            }
        }

        return hr;
    }
    
    HRESULT CopyTextureToSystemMemory(
        __in ID3D11Device* d3d11Device,
        __in ID3D11DeviceContext* d3d11Context,
        __in ID3D11Resource* d3d11Resource,
        __deref_out ScratchImage* image
        )
    {
        E_POINTER_Return(d3d11Device);
        E_POINTER_Return(d3d11Context);
        E_POINTER_Return(d3d11Resource);
        E_POINTER_Return(image);

        D3D11_RESOURCE_DIMENSION resType = D3D11_RESOURCE_DIMENSION_UNKNOWN;
        d3d11Resource->GetType(&resType);

        TRACEHR hr = S_OK;
        switch (resType)
        {
        case D3D11_RESOURCE_DIMENSION_TEXTURE1D:
            {
                Auto<ID3D11Texture1D> texture;
                
                hr = d3d11Resource->QueryInterface(__uuidof(ID3D11Texture1D), (void**)&texture);
                if (SUCCEEDED(hr))
                {
                    D3D11_TEXTURE1D_DESC desc;
                    texture->GetDesc(&desc);
                    
                    desc.BindFlags = 0;
                    desc.MiscFlags = 0;
                    desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
                    desc.Usage = D3D11_USAGE_STAGING;
                    
                    Auto<ID3D11Texture1D> stagingTexture;
                    hr = d3d11Device->CreateTexture1D(&desc, 0, &stagingTexture);
                    if (SUCCEEDED(hr))
                    {
                        d3d11Context->CopyResource(stagingTexture, d3d11Resource);
                        
                        TexMetadata mdata;
                        mdata.width = desc.Width;
                        mdata.height = mdata.depth = 1;
                        mdata.arraySize = desc.ArraySize;
                        mdata.mipLevels = desc.MipLevels;
                        mdata.miscFlags = 0;
                        mdata.format = desc.Format;
                        mdata.dimension = TEX_DIMENSION_TEXTURE1D;
                        
                        hr = image->Initialize(mdata);
                        if (SUCCEEDED(hr))
                        {
                            hr = _Capture(d3d11Context, stagingTexture, mdata, image);
                        }
                    }
                }
            }
            break;
        
        case D3D11_RESOURCE_DIMENSION_TEXTURE2D:
            {
                Auto<ID3D11Texture2D> texture;

                hr = d3d11Resource->QueryInterface(__uuidof(ID3D11Texture2D), (void**)&texture);
                if (SUCCEEDED(hr))
                {
                    D3D11_TEXTURE2D_DESC desc;
                    texture->GetDesc(&desc);
                    
                    Auto<ID3D11Texture2D> stagingTexture;
                    if (desc.SampleDesc.Count > 1)
                    {
                        desc.SampleDesc.Count = 1;
                        desc.SampleDesc.Quality = 0;
                        
                        Auto<ID3D11Texture2D> tmp;
                        hr = d3d11Device->CreateTexture2D(&desc, 0, &tmp);
                        if (SUCCEEDED(hr))
                        {
                            for (uint32 item = 0; item < desc.ArraySize; item++)
                            {
                                for (uint32 level = 0; level < desc.MipLevels; level++)
                                {
                                    uint32 index = D3D11CalcSubresource(level, item, desc.MipLevels);
                                    d3d11Context->ResolveSubresource(tmp, index, d3d11Resource, index, desc.Format);
                                }
                            }
                            
                            desc.BindFlags = 0;
                            desc.MiscFlags &= D3D11_RESOURCE_MISC_TEXTURECUBE;
                            desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
                            desc.Usage = D3D11_USAGE_STAGING;
                            
                            hr = d3d11Device->CreateTexture2D(&desc, 0, &stagingTexture);
                            if (SUCCEEDED(hr))
                            {
                                d3d11Context->CopyResource(stagingTexture, tmp);
                            }
                        }
                    }
                    else
                    {
                        desc.BindFlags = 0;
                        desc.MiscFlags &= D3D11_RESOURCE_MISC_TEXTURECUBE;
                        desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
                        desc.Usage = D3D11_USAGE_STAGING;
                        
                        hr = d3d11Device->CreateTexture2D(&desc, 0, &stagingTexture);
                        if (SUCCEEDED(hr))
                        {
                            d3d11Context->CopyResource(stagingTexture, d3d11Resource);
                        }
                    }

                    if (SUCCEEDED(hr))
                    {
                        TexMetadata mdata;
                        mdata.width = desc.Width;
                        mdata.height = desc.Height;
                        mdata.depth = 1;
                        mdata.arraySize = desc.ArraySize;
                        mdata.mipLevels = desc.MipLevels;
                        mdata.miscFlags = (desc.MiscFlags & D3D11_RESOURCE_MISC_TEXTURECUBE) ? TEX_MISC_TEXTURECUBE : 0;
                        mdata.format = desc.Format;
                        mdata.dimension = TEX_DIMENSION_TEXTURE2D;
                        
                        hr = image->Initialize(mdata);
                        if (SUCCEEDED(hr))
                        {
                            hr = _Capture(d3d11Context, stagingTexture, mdata, image);
                        }
                    }
                }
            }
            break;
        
        case D3D11_RESOURCE_DIMENSION_TEXTURE3D:
            {
                Auto<ID3D11Texture3D> texture;
                
                hr = d3d11Resource->QueryInterface(__uuidof(ID3D11Texture3D), (void**)&texture);
                if (SUCCEEDED(hr))
                {
                    D3D11_TEXTURE3D_DESC desc;
                    texture->GetDesc(&desc);
                    
                    desc.BindFlags = 0;
                    desc.MiscFlags = 0;
                    desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
                    desc.Usage = D3D11_USAGE_STAGING;
                    
                    Auto<ID3D11Texture3D> stagingTexture;
                    hr = d3d11Device->CreateTexture3D(&desc, 0, &stagingTexture);
                    if (SUCCEEDED(hr))
                    {
                        d3d11Context->CopyResource(stagingTexture, d3d11Resource);
                        if (SUCCEEDED(hr))
                        {
                            TexMetadata mdata;
                            mdata.width = desc.Width;
                            mdata.height = desc.Height;
                            mdata.depth = desc.Depth;
                            mdata.arraySize = 1;
                            mdata.mipLevels = desc.MipLevels;
                            mdata.miscFlags = 0;
                            mdata.format = desc.Format;
                            mdata.dimension = TEX_DIMENSION_TEXTURE3D;

                            hr = image->Initialize(mdata);
                            if (SUCCEEDED(hr))
                            {
                                hr = _Capture(d3d11Context, stagingTexture, mdata, image);
                            }
                        }
                    }
                }
            }
            break;
        
        default:
            hr = E_FAIL;
            break;
        }
        
        return hr;
    }

} // end DirectX namespace
